// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/dwio/common/output_stream.h>

namespace kumo::pollux::dwio::common {
    bool BufferedOutputStream::Next(
        void **buffer,
        int32_t *size,
        uint64_t increment) {
        // Try resize the buffer, if failed, flush it to make room
        if (!tryResize(buffer, size, 0, increment)) {
            flushAndReset(
                buffer, size, 0, {melon::Range(buffer_.data(), buffer_.size())});
        }
        return true;
    }

    void BufferedOutputStream::BackUp(int32_t count) {
        if (count > 0) {
            uint64_t unsignedCount = static_cast<uint64_t>(count);
            POLLUX_CHECK_LE(unsignedCount, buffer_.size(), "Can't backup that much!");
            buffer_.resize(buffer_.size() - unsignedCount);
        }
    }

    uint64_t BufferedOutputStream::flush() {
        const uint64_t dataSize = buffer_.size();
        if (dataSize > 0) {
            bufferHolder_.take(melon::Range(buffer_.data(), buffer_.size()));
            // resize to 0 is needed to make sure size() returns correct value
            buffer_.resize(0);
        }
        return dataSize;
    }

    void AppendOnlyBufferedStream::write(const char *data, size_t size) {
        while (size > 0) {
            if (MELON_UNLIKELY(bufferOffset_ == bufferLength_)) {
                POLLUX_CHECK(
                    outStream_->Next(
                        reinterpret_cast<void**>(&buffer_), &bufferLength_, size),
                    "Failed to allocate buffer.");
                bufferOffset_ = 0;
            }
            const size_t len =
                    std::min(static_cast<size_t>(bufferLength_ - bufferOffset_), size);
            ::memcpy(buffer_ + bufferOffset_, data, len);
            bufferOffset_ += static_cast<int32_t>(len);
            data += len;
            size -= len;
        }
    }

    uint64_t AppendOnlyBufferedStream::flush() {
        outStream_->BackUp(bufferLength_ - bufferOffset_);
        bufferOffset_ = 0;
        bufferLength_ = 0;
        buffer_ = nullptr;
        return outStream_->flush();
    }
} // namespace kumo::pollux::dwio::common
